import os.path
import time
import pandas as pd
import numpy as np
from scipy.stats import rankdata
# from numba import njit
import pyarrow as pa
import pyarrow.compute as pc
# from cython_optimized_functions import stddev_cython_incremental_calculation
import bottleneck as bn

def stddev(x, window=10):
    print(f"stdddev use {stddev_bottleneck}")
    return stddev_bottleneck(x, window)

def stddev_numpy(x, window=10):
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window,axis = 0)
    return np.append(np.full([window-1,np.size(x, 1)],np.nan),np.std(a_rolled, axis=-1, ddof=1),axis = 0)

def stddev_numpy_optimized(x, window=10):
    """
    原先的 stddev 优化版本
    先创建矩阵, 避免复制

    总耗时 0.250 左右
    没啥太明显的性能提升

    :param x:
    :param window:
    :return:
    """
    # 创建滑动窗口视图
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window, axis=0)

    # 预分配结果矩阵，前 window-1 行填充 NaN
    result = np.empty((x.shape[0], x.shape[1]))
    result[:window-1, :] = np.nan

    # 计算标准差，并将结果放入预分配的矩阵中
    result[window-1:, :] = np.std(a_rolled, axis=-1, ddof=1)

    return result

# @njit
# def stddev_numba(x, window=10):
#     """
#     使用 numba 加速
#     总耗时 2.8187830448150635

#     结论:
#     加不加 @njit 效果一样
#     numba 不稳定, 弃用
#     :param x:
#     :param window:
#     :return:
#     """
#     rows, cols = x.shape
#     result = np.full((rows, cols), np.nan)

#     for col in range(cols):
#         # 计算第一个窗口的平均值和平方和
#         window_sum = np.sum(x[:window, col])
#         window_sq_sum = np.sum(x[:window, col]**2)

#         for i in range(window - 1, rows):
#             if i >= window:
#                 window_sum -= x[i - window, col]
#                 window_sq_sum -= x[i - window, col]**2

#             window_sum += x[i, col]
#             window_sq_sum += x[i, col]**2

#             mean = window_sum / window
#             variance = (window_sq_sum / window) - mean**2
#             result[i, col] = np.sqrt(variance)

#     return result

def stddev_bottleneck(x, window=10):
    return bn.move_std(x, window=window)


def rank(x):
    print(f"rank use {rank_bottleneck}")
    return rank_bottleneck(x)

def rank_scipy(x):
    return rankdata(x,method='min',axis=1)/np.size(x, 1)

def rank_bottleneck(x):
    return bn.rankdata(x, axis=1)/x.shape[1]

def pct_change(ndarr):
    pct_change = (np.diff(ndarr, axis=0) / ndarr[:-1])
    return np.vstack([np.full(ndarr.shape[1], np.nan), pct_change])

def ts_argmax(x, window=10):
    print(f"ts_argmax use {ts_argmax_bottleneck}")
    return ts_argmax_bottleneck(x, window)

def ts_argmax_numpy(x, window=10):
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window,axis = 0)
    return np.append(np.full([window-1,np.size(x, 1)],np.nan),np.argmax(a_rolled, axis=-1) + 1,axis = 0)

def ts_argmax_bottleneck(x, window=10):
    return bn.move_argmax(x, window)


